#!/usr/bin/env python3.9
# coding=utf-8

"""
Copyright (c) Huawei Technologies Co., Ltd. 2020-2022. All rights reserved.
Description: python api test.
Author: MindX SDK
Create: 2022
History: NA
"""

import os
import logging
import time
from multiprocessing import Process, JoinableQueue, Queue, Manager, Pool
import traceback
import numpy as np
import cv2
from sklearn import preprocessing
from mindx.sdk import base
from mindx.sdk.base import Image, ImageProcessor, Size, Tensor, Dim, Color
from mindx.sdk.base import VideoDecoder, VideoDecodeConfig, VdecCallBacker

import body.main_body_batch_v3 as Body
import portrait.main_face_batch_v3 as Face

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)
device_id = 0
channel_id = 0
SKIP_INTERVAL = 3
BATCH_SIZE=16
VIDEO_HEIGHT = 4096
VIDEO_WIDTH = 2160
VIDEO_BREAK_FRAME = 100
JPG_PATH = "./data/img_files"

def process0_video_decode(car_union_q:Queue, carplate_q:Queue, body_q:Queue, face_q: Queue, video_stream_path: str, channelId: int):
    """
    decode_q_pro: Queue, decoded frames to be put in
    video_stream_path: str, video to decode from
    channelId: int, video decode channel
    """
    try:
        logger.info("=========================== Start process0_video_decode ===========================")
        image_processor = ImageProcessor(device_id)  # initialize mxbase image_process

        def vdec_callback(decodeImage, channelId, frameId):
            """
            Define videoDecoder callback function
            """
            logger.debug(
                f"process0_video_decode decodeImage channelId: {channelId}, frameId: {frameId} width: {decodeImage.width}, height: {decodeImage.height}")

            # 1. Calculate preprocess resize and padding config
            decodeImage.to_host()
            image_ori = np.array(decodeImage.to_tensor())  # NHWC, RGB

            car_union_q.put((frameId, image_ori, decodeImage.original_height, decodeImage.original_width))
            carplate_q.put((frameId, image_ori, decodeImage.original_height, decodeImage.original_width))
            body_q.put((frameId, image_ori, decodeImage.original_height, decodeImage.original_width))
            face_q.put((frameId, image_ori, decodeImage.original_height, decodeImage.original_width))

            logger.debug("process0_video_decode producer queue, elements {}".format(decode_q_pro.qsize()))

        # 1. initialize video decoder callback
        vdec_callbacker = VdecCallBacker()
        vdec_callbacker.registerVdecCallBack(vdec_callback)

        # 2. initialize video decoder config
        vdec_config = VideoDecodeConfig()
        vdec_config.skipInterval = SKIP_INTERVAL
        vdec_config.outputImageFormat = base.rgb
        vdec_config.width = VIDEO_WIDTH
        vdec_config.height = VIDEO_HEIGHT

        # 3. initialize video decode object
        video_decoder = VideoDecoder(vdec_config, vdec_callbacker, deviceId=device_id, channelId=channelId)

        # 4. use ffmpeg av and VideoDecoder to decode video
        import av
        with av.open(video_stream_path) as container:
            logger.debug(
                f"============================process0_video_decode start decode stream =========================")
            for frame_id, packet in enumerate(container.demux()):
                logger.debug(f"process0_video_decode curr frame id: {frame_id}, size: {packet.size}")
                if packet.size == 0:
                    break
                try:
                    video_decoder.decode(packet, frame_id)
                except Exception as err:
                    logger.error("process0_video_decode decode error: ", err)
                if frame_id == VIDEO_BREAK_FRAME:
                    break
        logger.info("=============== Finish process0_video_decode ========================")
        del video_decoder
    except Exception as e:
        print("process0_video_decode failed.", repr(e))
        traceback.print_exc()

def process0_read_image(car_union_q:Queue, carplate_q:Queue, body_q:Queue, face_q: Queue, pid):
    try:
        logger.info("=========================== Start process0_read_image ===========================")
        image_processor = ImageProcessor(device_id)  # initialise mxbase image_process
        count = 0

        for img_file in os.listdir(JPG_PATH):
            logger.info(f"process0_read_image{img_file}")
            decodeImage = image_processor.decode(os.path.join(JPG_PATH, img_file), base.rgb)  # NHWC
            # 1. Calculate preprocess resize and padding config
            decodeImage.to_host()
            image_ori = np.array(decodeImage.to_tensor())  # NHWC, RGB

            car_union_q.put((img_file, image_ori, decodeImage.original_height, decodeImage.original_width))
            carplate_q.put((img_file, image_ori, decodeImage.original_height, decodeImage.original_width))
            body_q.put((img_file, image_ori, decodeImage.original_height, decodeImage.original_width))
            face_q.put((img_file, image_ori, decodeImage.original_height, decodeImage.original_width))
            count += 1

        logger.info(f"process0_read_image pid:{pid} finished, put {count} images")
    except Exception as e:
        print("process0_read_image failed.", repr(e))
        traceback.print_exc()

if __name__ == "__main__":
    start = time.time()
    base.mx_init()

    NUM_DECODE_PROCESS = 20
    NUM_INFER_PROCESS = 4
    NUM_PATTERN_RECOG_PROCESS = 4
    NUM_REID_PROCESS = 4
    manager = Manager()
    q_decode_car_union = manager.Queue()
    q_decode_carplate = manager.Queue()

    q_decode_body = manager.Queue()
    q_rec_body = manager.Queue()
    q_reid_body = manager.Queue()

    q_decode_face = manager.Queue()
    q_rec_face = manager.Queue()

    pool = Pool()

    for i in range(NUM_DECODE_PROCESS):
        # pool.apply_async(process0_video_decode, args=(q_decode, VIDEO_PATH, i))
        pool.apply_async(process0_read_image, args=(q_decode_car_union, q_decode_carplate, q_decode_body, q_decode_face, i,))
    for i in range(NUM_INFER_PROCESS):
        pool.apply_async(Body.process1_infer, args=(q_decode_body, q_rec_body, q_reid_body, BATCH_SIZE, i))
    for i in range(NUM_PATTERN_RECOG_PROCESS):
        pool.apply_async(Body.process2_pattern_rec, args=(q_rec_body, i,))
    for i in range(NUM_REID_PROCESS):
        pool.apply_async(Body.process3_reid, args=(q_reid_body, i,))

    for i in range(NUM_INFER_PROCESS):
        pool.apply_async(Face.process1_infer, args=(q_decode_face, q_rec_face, BATCH_SIZE, i))
    for i in range(NUM_PATTERN_RECOG_PROCESS):
        pool.apply_async(Face.process2_pattern_rec, args=(q_rec_face, i))

    pool.close()
    pool.join()

    end = time.time()
    print("total time", end - start)
    print("avg time", (end - start) / (NUM_DECODE_PROCESS * len(os.listdir(JPG_PATH))))
